open Variable
open Type
open Object
open List

let rec t1_to_t2 e =
  match e with
    Var(name) ->  T2_Var(name)
  | Binary(e1, op, e2) -> T2_Binary(t1_to_t2 e1, op, t1_to_t2 e2)
  | Unary(op, e) -> T2_Unary(op, t1_to_t2 e)
  | Assignment(var, e) -> T2_Assignment(t1_to_t2 var, t1_to_t2 e)
  | Int(v) -> T2_Int(v)
  | Float(v) -> T2_Float(v)
  | Bool(v) -> T2_Bool(v)
  | String(v) -> T2_String(v)
  | Array(ls) -> T2_Array(List.map t1_to_t2 ls)
  | Array_index(_,_) -> error "todo t1_to_t2"
  | Struct_If(if_list) -> T2_Struct_If(List.map t1_to_t2_if if_list )
  | Print(var_list) -> T2_Print(List.map t1_to_t2 var_list)
  | Return(e) -> T2_Return(t1_to_t2 e)
  | Function(name, args) -> T2_Function(name, List.map t1_to_t2 args)
  | For(var, b, e, body) ->
    let new_var = t1_to_t2 var
    and new_b = t1_to_t2 b
    and new_e = t1_to_t2 e
    and new_body = List.map t1_to_t2 body in
    T2_For(new_var, new_b, new_e, new_body)
  | While(cond, body) ->
    let new_cond = t1_to_t2 cond
    and new_body = List.map t1_to_t2 body in
    T2_While(new_cond, new_body)
  | Decl_Function(name, args, body) ->
      T2_Decl_Function(name, args, List.map t1_to_t2 body)

and t1_to_t2_if expr =
  match expr with
    If(cond, expr_list) -> let new_cond = t1_to_t2 cond and
               new_expr = (List.map t1_to_t2 expr_list) in
              (T2_If(new_cond, new_expr))
  | Else(expr_list) -> let new_expr = List.map t1_to_t2 expr_list in
              (T2_Else(new_expr))

;;


(* change variable name is expr, should be called with List.map ... *)
let rec rename_var old_name new_name expr =
  match expr with
    Var(name) ->  (if (String.compare name old_name == 0)
          then begin print_string (name^" > "^new_name^"\n"); (Var(new_name))
          end else expr)
  | Binary(e1, op, e2) -> let new_e1 = rename_var old_name new_name e1
              and new_e2 = rename_var old_name new_name e2 in
              (Binary(new_e1, op, new_e2))
  | Unary(op, e) -> (Unary(op, rename_var old_name new_name e))
  | Assignment(var, e) -> let new_e = rename_var old_name new_name e in
              (Assignment(rename_var old_name new_name var, new_e))
  | Struct_If(if_list) ->
      (Struct_If(List.map (rename_var_if old_name new_name) if_list ))
  | Print(var_list) -> Print(List.map (rename_var old_name new_name ) var_list)
  | Return(e) -> Return(rename_var old_name new_name e)
  | Function(name, args) ->
      Function(name, List.map (rename_var old_name new_name) args)
  | For(var, b, e, body) ->
    let new_var = rename_var old_name new_name var
    and new_b = rename_var old_name new_name b
    and new_e = rename_var old_name new_name e
    and new_body = List.map (rename_var old_name new_name) body in
    For(new_var, new_b, new_e, new_body)
  | While(cond, body) ->
    let new_cond = rename_var old_name new_name cond
    and new_body = List.map (rename_var old_name new_name) body in
    While(new_cond, new_body)
  | Decl_Function(name, args, body) -> Decl_Function(name, args, body)
  | _ -> expr
  

and rename_var_if old_name new_name expr =
  match expr with
    If(cond, expr_list) ->
      let new_cond = rename_var old_name new_name cond and
      new_expr = (List.map (rename_var old_name new_name) expr_list) in
        (If(new_cond, new_expr))
  | Else(expr_list) ->
      let new_expr = List.map (rename_var old_name new_name) expr_list in
        (Else(new_expr))
;;

(* functions to move declaration out of if *)
let rec move_var_decl expr_list =
    match expr_list with
      [] -> []
    | head::tail -> 
      begin match head with
        Assignment(var,value) ->
        let name = get_name_t1 var in
        if ( is_present name ) 
          then (move_var_decl tail)
        else begin
          add_var TNil name ;
          T2_Decl_Var(name)::(move_var_decl tail)
        end
      | _ -> move_var_decl tail
      end
;;

let rec move_decl_if if_list =
  match if_list with
    [] -> []
  | head::tail ->
    begin match head with
      If(cond,expr_list) -> let x = move_var_decl expr_list in
                  x@(move_decl_if tail)
    | Else(expr_list) -> let x = move_var_decl expr_list in
                x@(move_decl_if tail)
    end
;;

(* split declaration and assignment + move decl_var out of if
+ add decl_var before for loop *)
let rec split_decl ast =
  match ast with
      [] -> []
    | head::tail -> 
      begin match head with
        Struct_If(if_list) -> let bef = move_decl_if if_list in
        (bef@[t1_to_t2 head]@(split_decl tail))
      | Assignment(var,value) ->
        let name = get_name_t1 var in
        let new_ass = T2_Assignment(t1_to_t2 var, t1_to_t2 value) in
        if ( is_present name)
          then new_ass::(split_decl tail)
        else begin
          add_var TNil name;
          T2_Decl_Var(name)::new_ass::(split_decl tail)
        end
      | For(var, b, e, body) ->
        (* rename loop variable to make sure it has an non ObjVar type *)
        let name = get_name_t1 var in
        let new_name = get_new_name () in
        (* add an internal var which is initialized at i
        to reporduce ruby behaviour *)
        let n_name2 = get_new_name () in
        let new_body = Assignment(Var(n_name2), Var(new_name))::body in
        let new_body2 = (List.map (rename_var name n_name2 ) new_body) in
        let bef_for = move_var_decl body in
        let new_for = bef_for@[T2_Decl_Var(n_name2)]
          @[T2_For(T2_Var(new_name), t1_to_t2 b, t1_to_t2 e, List.map t1_to_t2 new_body2)] and
        ass_after = T2_Assignment(t1_to_t2 var, T2_Var(n_name2)) in
        add_var TNil new_name;
        add_var TNil n_name2;
        if ( not (is_present name))
          then begin
          add_var TNil name;
          T2_Decl_Var(name)::T2_Decl_Var(new_name)::new_for
            @[ass_after]@(split_decl tail)
        end else
          T2_Decl_Var(new_name)::new_for@[ass_after]@(split_decl tail)

      | While(cond, body) ->
        let bef_w = move_var_decl body in
        bef_w@[t1_to_t2 head]@(split_decl tail)
      | Decl_Function(name, args, body) ->
        let old_env = !env in
        env := [];
        (* rename all args to avoid name conflicts *)
        let new_args=List.map (function name -> (name, get_new_name())) args in
        let rec get_new_body args_list body =
          match args_list with
            [] -> body
          | (name, new_name)::tail ->
              get_new_body tail (List.map (rename_var name new_name) body)
        in
        let new_body = get_new_body new_args body in
        (* add all args in env *)
        let new_args2 = List.map (function x -> snd(x)) new_args in
        List.iter (function x -> (add_var TNil x)) new_args2;

        let new_body2 = List.map split_decl [new_body] in
        env := old_env;
        T2_Decl_Function(name, new_args2, hd new_body2)::(split_decl tail)
      | Return(e) -> (* return will always be a variable *)
        let new_e = begin match e with
                      Var(name) -> [T2_Return(T2_Var(name))]
                    | Int(v) -> [T2_Return(T2_Int(v))]
                    | Float(v) -> [T2_Return(T2_Float(v))]
                    | String(v) -> [T2_Return(T2_String(v))]
                    | Array(v) -> [T2_Return(T2_Array(map t1_to_t2 v))]
                    | Array_index(name, idx) ->
                        [T2_Return(T2_Array_index(name, t1_to_t2 idx))]
                    | _ -> let new_name = get_new_name () in
                           [T2_Decl_Var(new_name);T2_Assignment(T2_Var(new_name), t1_to_t2 e);
                           T2_Return(T2_Var(new_name))]
                    end
                    in
        new_e@(split_decl tail)
        
      | x -> (t1_to_t2 x)::(split_decl tail)
      end
;;


let rec move_decl ast =
  let rec move_decl_int ast fun_decl other =
    match ast with
      [] -> fun_decl@other
    | head::tail -> 
      begin match head with
        T2_Decl_Function(_, _, _) -> move_decl_int tail (head::fun_decl) other
      | _ -> move_decl_int tail fun_decl (head::other)
      end
  in move_decl_int ast [] []
;;

let move = function ast ->
  let new_ast = split_decl ast in move_decl ((List.rev)new_ast) ;;
